import numpy as np
import ncnn
import torch
import time
import matplotlib.pyplot as plt

def test_inference(num_warmup=50, num_timed=100):
    torch.manual_seed(0)
    in0 = torch.rand(1, 3, 640, 640, dtype=torch.float)

    # ncnn.create_gpu_instance()
    print("get gpu count: ", ncnn.get_gpu_count())

    # 加载模型（仅一次）
    net = ncnn.Net()
    net.opt.use_vulkan_compute = True
    net.set_vulkan_device(0)
    net.load_param("yolo11n_ncnn_model/model.ncnn.param")
    net.load_model("yolo11n_ncnn_model/model.ncnn.bin")

    # 预热阶段
    print(f"Start warmup ({num_warmup} iterations)...")
    with net.create_extractor() as ex:
        for _ in range(num_warmup):
            ex.input("in0", ncnn.Mat(in0.squeeze(0).numpy()).clone())
            _, _ = ex.extract("out0")
    print("Warmup done.\n")

    # 正式推理并记录时间
    print(f"Start timing inference ({num_timed} iterations)...")
    timings = []
    with net.create_extractor() as ex:
        for _ in range(num_timed):
            # 计时
            start_time = time.perf_counter()
            # 输入数据准备
            ex.input("in0", ncnn.Mat(in0.squeeze(0).numpy()).clone())
            # 推理
            _, out0 = ex.extract("out0")
            end_time = time.perf_counter()

            timings.append((end_time - start_time) * 1000)  # 转换为毫秒

    print("Inference timing completed.\n")

    # 计算平均时间和标准差
    avg_time = np.mean(timings)
    std_time = np.std(timings)
    min_time = np.min(timings)
    max_time = np.max(timings)
    print(f"Average Inference Time: {avg_time:.2f} ms")
    print(f"Standard Deviation: {std_time:.2f} ms")
    print(f"Min / Max: {min_time:.2f} ~ {max_time:.2f} ms")

    # 绘制直方图
    plt.figure(figsize=(10, 6))
    plt.hist(timings, bins=20, color='skyblue', edgecolor='black')
    plt.title("NCNN Inference Latency Distribution (ms)")
    plt.xlabel("Latency (ms)")
    plt.ylabel("Frequency")
    plt.axvline(avg_time, color='red', linestyle='dashed', linewidth=2, label=f'Mean = {avg_time:.2f} ms')
    plt.legend()
    plt.grid(True)
    plt.show()

    return avg_time, std_time, timings


if __name__ == "__main__":
    avg_time, std_time, timings = test_inference(num_warmup=5000, num_timed=10000)